import pandas as pd
import os,csv
import argparse

def add_LD(a):
    VarPUnPLP_related_List = resultFinly[resultFinly["UnPLP"]==a]
    VarPUnPLP_related_index_List = VarPUnPLP_related_List.index.tolist()
    
    if VarPUnPLP_related_index_List:
        r_squareds, others = [], []
        for idnex_ in VarPUnPLP_related_index_List:
            r_squared = str(resultFinly.loc[idnex_, "r_squared"])
            
            total = resultFinly.loc[idnex_, " total"]
            Sum_A = resultFinly.loc[idnex_, "Sum_A"]
            Sum_B = resultFinly.loc[idnex_, "Sum_B"]
            Sum_AB = resultFinly.loc[idnex_, "Sum_AB"]
            P_B = resultFinly.loc[idnex_, "VarP_A"]
            P_A = resultFinly.loc[idnex_, "VarP_B"]
            other = P_A + ":" + P_B + ":" + str(total) + ":" + str(Sum_A) + ":" + str(Sum_B) + ":" + str(Sum_AB)
            r_squareds.append(r_squared)
            others.append(other)
    else:
        return "-", "-"
    return  "|".join(r_squareds), "|".join(others)

def add_MQ(a):
    if a in tag1_MQ.keys():
        return tag1_MQ[a]
    else:
        return "-"
    
def join_tag(df):
    return df.iloc[0] + "_" + str(df.iloc[1]) + "_" + df.iloc[3] + "_" + df.iloc[4]
def join_tag1(df):
    return df.iloc[0] + "_" + str(df.iloc[1])

# LD
outputPath="/lustre/home/acct-medfzx/medfzx-lkw/project/pangenome/other/CYP21A2_3rd/cah_hgrc.GT.rmCHM13.LDtest.r02.PLP.csv"
resultFinly = pd.read_csv(outputPath, sep=",", index_col=0)
resultFinly = resultFinly.reset_index(drop=True)
# MQ
# VCF annovar


def add_LD_MQ(NGS_annovar_path):
    df_clinvar=pd.read_csv(NGS_annovar_path,sep="\t")
    # 移除Otherinfo1	Otherinfo2	Otherinfo3	Otherinfo4	Otherinfo5	Otherinfo6	Otherinfo7	Otherinfo8	Otherinfo11行
    df_clinvar = df_clinvar.drop(['Otherinfo1', 'Otherinfo2', 'Otherinfo3', 'Otherinfo4', 'Otherinfo5', 'Otherinfo6', 'Otherinfo7', 'Otherinfo8', 'Otherinfo11'], axis=1)
    df_clinvar['tag'] = df_clinvar.apply(join_tag, axis=1)
    df_clinvar['tag1'] = df_clinvar.apply(join_tag1, axis=1)

    df_clinvar[['r_squared', 'other']] = df_clinvar.apply(lambda x: pd.Series(add_LD(x['tag'])), axis=1)
    df_clinvar['MQ'] = df_clinvar["tag1"].apply(lambda x: add_MQ(x))
    df_clinvar = df_clinvar.drop(['tag', 'tag1'], axis=1)
    return df_clinvar


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="get average Mapping Quality at a position.")
    parser.add_argument('-i', "--input", help="Using mpileup get the txt file from bam.")
    parser.add_argument('-m', "--mqcsv", help="MQ file path.", default="./")
    # parser.add_argument('-n', "--number", type=int, help="The number of generate sequence.", default="10")
    args = parser.parse_args()
    df_MQ=pd.read_csv(args.mqcsv, quoting=csv.QUOTE_NONE, sep='\t', names=["chr", "start", "end", "MQ"])
    df_MQ['tag1'] = df_MQ.apply(join_tag1, axis=1)
    df_MQ.head(5)
    tag1_MQ = {key: value for key, value in zip(df_MQ['tag1'], df_MQ['MQ'])}
    resultDf=add_LD_MQ(args.input)
    resultDf.to_csv(args.input.split(".")[0] + ".csv")
    
    
    
